import network
import socket
import machine
import time
from machine import I2S, Pin

# ================= 配置区域 =================
SSID = "Astro_RDX"       
PASSWORD = "17060840"    
PORT = 50005

# 引脚
PIN_SCK = 13     # BCLK
PIN_WS = 14      # LRCK
PIN_SD_SPK = 15  # DIN

# 音频参数
SAMPLE_RATE = 16000
# 关键点：使用 16位，不要用 32位，减少数据量和计算量
SAMPLE_BITS = 16 
# ==========================================

led = machine.Pin("LED", machine.Pin.OUT)

def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    # 关闭节能模式 (关键!)
    wlan.config(pm=0xa11140)
    
    if not wlan.isconnected():
        print(f"Connecting to {SSID}...")
        wlan.connect(SSID, PASSWORD)
        while not wlan.isconnected():
            time.sleep(0.5)
            print(".", end="")
    
    ip = wlan.ifconfig()[0]
    led.on()
    print(f"\nIP: {ip}")
    return ip

def main():
    ip = connect_wifi()
    if not ip: return

    print(f"Init I2S: Rate={SAMPLE_RATE}, Bits={SAMPLE_BITS}, Mode=MONO")
    
    # === 关键修改 ===
    # 1. format=I2S.MONO (解决夹子音)
    # 2. bits=16 (解决 Python 运算慢导致的卡顿)
    audio_out = I2S(1,
        sck=Pin(PIN_SCK), ws=Pin(PIN_WS), sd=Pin(PIN_SD_SPK),
        mode=I2S.TX,
        bits=SAMPLE_BITS,
        format=I2S.MONO, 
        rate=SAMPLE_RATE,
        ibuf=40000  # 大缓冲区，防止没数据时卡顿
    )

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(('0.0.0.0', PORT))
    sock.settimeout(0.1) # 超时设短一点

    print(f"Listening on {PORT}...")
    
    # 预热缓冲区：静音防爆音
    silence = bytearray(1024)
    for _ in range(5):
        audio_out.write(silence)

    try:
        while True:
            try:
                # 接收数据
                data, addr = sock.recvfrom(1024)
                
                if data and len(data) > 0:
                    if data == b'END_OF_STREAM':
                        print("Stream End")
                        # 播放一点静音防止结尾爆音
                        audio_out.write(silence)
                    else:
                        # === 0 计算量直接转发 ===
                        # 只要格式对齐 (16bit mono)，直接 write 最快
                        audio_out.write(data)
                        
            except OSError:
                pass
                
    except KeyboardInterrupt:
        pass
    finally:
        audio_out.deinit()
        sock.close()
        led.off()

if __name__ == "__main__":
    main()